import time
import numpy as np
import mujoco
import mujoco.viewer

import modern_robotics as mr

from src.utils import TransformationUtils
from src.robot import Robot, Frankie
from src.wheel_controller import WheelController, DifferentialDriveWheelController
from src.motion_controller import BaseController, HighLevelController, ArmController, RedundancyResolutionController, \
    GripperController

if __name__ == '__main__':
    model = mujoco.MjModel.from_xml_path("../assets/franka_emika_panda/scene.xml")
    data = mujoco.MjData(model)
    robot: Robot = Frankie()
    robot.q = [0.0, 0.0, 0.0, -0.3, 0.0, -2.2, 0.0, 2.0, np.pi / 4]

    joint1_joint_name = "joint1"
    joint1_joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint1_joint_name)

    box_joint1_name = "box1_joint"
    box_joint1_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, box_joint1_name)
    box_joint2_name = "box2_joint"
    box_joint2_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, box_joint2_name)

    q0 = np.zeros_like(data.qpos)
    q0[3] = 1.0
    q0[joint1_joint_id + 1 + 12] = -0.3
    q0[joint1_joint_id + 3 + 12] = -2.2
    q0[joint1_joint_id + 5 + 12] = 2.0
    q0[joint1_joint_id + 6 + 12] = np.pi / 4
    q0[box_joint1_id + 12] = 2.0
    q0[box_joint1_id + 12 + 1] = 1.0
    q0[box_joint1_id + 12 + 2] = 0.6
    q0[box_joint1_id + 12 + 3] = 1
    q0[box_joint2_id + 18] = 6.0
    q0[box_joint2_id + 18 + 1] = 2.0
    q0[box_joint2_id + 18 + 2] = 0.6
    q0[box_joint2_id + 18 + 3] = 1

    ctrl0 = np.zeros_like(data.ctrl)
    ctrl0[2: 9] = q0[joint1_joint_id + 12: joint1_joint_id + 7 + 12]
    ctrl0[9] = 255
    mujoco.mj_setState(model, data, q0, mujoco.mjtState.mjSTATE_QPOS)
    mujoco.mj_setState(model, data, ctrl0, mujoco.mjtState.mjSTATE_CTRL)
    mujoco.mj_forward(model, data)

    gripper_site_name = 'gripper_site'
    gripper_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, gripper_site_name)
    R_gripper = data.site_xmat[gripper_site_id].reshape(3, -1)
    t_gripper = data.site_xpos[gripper_site_id]
    T_gripper = mr.RpToTrans(R_gripper, t_gripper)

    base_site_name = 'base_arm'
    base_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, base_site_name)
    R_base = data.site_xmat[base_site_id].reshape(3, -1)
    t_base = data.site_xpos[base_site_id]
    T_base = mr.RpToTrans(R_base, t_base)

    r = 0.05
    w = 0.2
    wheel_controller: WheelController = DifferentialDriveWheelController(r, w)

    object1_site_name = 'grasp_object1'
    object1_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, object1_site_name)
    R_object1 = data.site_xmat[object1_site_id].reshape(3, -1)
    t_object1 = data.site_xpos[object1_site_id]
    T_object1 = mr.RpToTrans(R_object1, t_object1)

    location1_site_name = 'drop_location1'
    location1_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, location1_site_name)
    T_location1 = mr.RpToTrans(data.site_xmat[location1_site_id].reshape(3, -1), data.site_xpos[location1_site_id])

    object2_site_name = 'grasp_object2'
    object2_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, object2_site_name)
    T_object2 = mr.RpToTrans(data.site_xmat[object2_site_id].reshape(3, -1), data.site_xpos[object2_site_id])

    location2_site_name = 'drop_location2'
    location2_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, location2_site_name)
    T_location2 = mr.RpToTrans(data.site_xmat[location2_site_id].reshape(3, -1), data.site_xpos[location2_site_id])

    object3_site_name = 'grasp_object3'
    object3_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, object3_site_name)
    T_object3 = mr.RpToTrans(data.site_xmat[object3_site_id].reshape(3, -1), data.site_xpos[object3_site_id])

    high_level_controller = HighLevelController()
    high_level_controller.add_point(T_object1)
    high_level_controller.add_point(T_location1)
    high_level_controller.add_point(T_object2)
    high_level_controller.add_point(T_location2)
    high_level_controller.add_point(T_object3)

    control_timestep = 0.01
    n_steps = round(control_timestep / model.opt.timestep)
    base_controller = BaseController()
    arm_controller = ArmController(control_timestep, robot)
    T_gb = mr.TransInv(T_base) @ T_gripper
    arm_controller.reset(T_gb)
    redundancy_resolution_controller = RedundancyResolutionController(robot)
    gripper_controller = GripperController(control_timestep)

    R_closest = np.eye(3)

    num = 0
    with mujoco.viewer.launch_passive(model, data) as viewer:
        while viewer.is_running():

            start_time = time.time()

            if num % n_steps == 0:
                R_base = data.site_xmat[base_site_id].reshape(3, -1)
                t_base = data.site_xpos[base_site_id]
                T_base = mr.RpToTrans(R_base, t_base)

                R_gripper = data.site_xmat[gripper_site_id].reshape(3, -1)
                t_gripper = data.site_xpos[gripper_site_id]
                T_gripper = mr.RpToTrans(R_gripper, t_gripper)

                T_target, T_next = high_level_controller.ctrl()
                succeed, v_base, time_in, T_closest = base_controller.ctrl(T_target, T_next, T_base, R_base)

                T_gb = mr.TransInv(T_base) @ T_gripper
                v_gripper_desired, v_base_desired = arm_controller.ctrl(T_target, time_in, v_base, T_base, T_gb,
                                                                        T_closest)
                qd = redundancy_resolution_controller.ctrl(v_gripper_desired, v_base_desired)

                if np.linalg.norm(T_target[:3, 3] - t_gripper) < 0.03:
                    if high_level_controller.current_id % 2 == 0:
                        arrive, gripper_ctrl = gripper_controller.close()
                        if arrive:
                            high_level_controller.update()
                            T_gb = mr.TransInv(T_base) @ T_gripper
                            arm_controller.reset(T_gb)
                    if high_level_controller.current_id % 2 == 1:
                        arrive, gripper_ctrl = gripper_controller.open()
                        if arrive:
                            high_level_controller.update()
                            T_gb = mr.TransInv(T_base) @ T_gripper
                            arm_controller.reset(T_gb)
                else:
                    gripper_ctrl = gripper_controller.get()

            q = robot.q
            q += qd * model.opt.timestep

            wheel_velocity = wheel_controller.ctrl(qd[1], 0.0, qd[0])
            data.ctrl[:2] = wheel_velocity
            data.ctrl[2:9] = q[2:9]
            data.ctrl[9] = gripper_ctrl
            mujoco.mj_step(model, data)

            R_base = data.site_xmat[base_site_id].reshape(3, -1)
            t_base = data.site_xpos[base_site_id]
            T_base = mr.RpToTrans(R_base, t_base)

            R_gripper = data.site_xmat[gripper_site_id].reshape(3, -1)
            t_gripper = data.site_xpos[gripper_site_id]
            T_gripper = mr.RpToTrans(R_gripper, t_gripper)

            robot.set_base(T_base)
            q[:2] = 0
            robot.q = q

            viewer.sync()

            num += 1

            end_time = time.time()
            delta_time = end_time - start_time
            if model.opt.timestep - delta_time > 0:
                time.sleep(model.opt.timestep - delta_time)
